# Copyright (C) 2021.Huawei Technologies Co., Ltd. All rights reserved.
# coding: utf-8
import configparser
import os
import shutil
import pickle
import time
from io import BytesIO

from moxing.framework import cloud_utils
from security.api import symmetric
from security.exceptions import DewError
from security.key.dew import Dew
from security.utils import hex_string_to_bytes
from security.wrappers.kmc_wrapper import CryptoAlgorithm, SignAlgorithm

CUR_RANK_ID = os.environ.get('RANK_ID')


def get_dew_config(dew_config_file):
    if not os.path.exists(dew_config_file):
        raise DewError(f"DEW config file not exists : {dew_config_file}")
    # 读取DEW配置文件内容
    _CONFIG_PARAMS = ('key_id', 'project_id', 'base_url', 'crypto', 'signature')
    conf_parser = configparser.RawConfigParser()
    conf_parser.read(dew_config_file)
    dew_config = {}
    if 'dew' not in conf_parser:
        raise DewError(f'no dew section in config: {dew_config_file}')
    section_conf = dict(conf_parser.items('dew'))
    for item in _CONFIG_PARAMS:
        if item not in section_conf:
            raise DewError(f'no {item} in config: {dew_config_file}')
        dew_config[item] = section_conf.get(item)
    if 'verify' in section_conf:
        dew_config['verify'] = conf_parser.getboolean('dew', 'verify')
    if 'cipher_key' in section_conf:
        dew_config['cipher_key'] = section_conf.get('cipher_key')
    auth_info = cloud_utils.get_auth()
    dew_config['ak'] = auth_info.AK
    dew_config['sk'] = auth_info.SK
    return dew_config


def decrypt_file_stream(src_file, dst_dir, plain_key):
    """
    Decrypt src_file and save result to dst_dir
    Args:
        src_file: cipher file to be decrypted
        dst_dir: the directory to save result file
        plain_key: the plain data key
    Returns:
        None
    """
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    file_name = os.path.basename(src_file)
    decrypted_data_file = os.path.join(dst_dir, file_name)
    print(f"[{CUR_RANK_ID}] decrypt {src_file} start ...")
    # 对文件进行流式解密
    dec_stream = symmetric.decrypt_stream(plain_key, src_file,
                                          symmetric.CHUNK_SIZE_1_MB)
    with open(decrypted_data_file, 'wb') as f:
        for dec_data in dec_stream:
            f.write(dec_data)
    print(f'[{CUR_RANK_ID}] decrypt {src_file} successfully')


def decrypt_dir(dew_config_file, decrypt_dir_name, user_dir_prefix='/cache',
                decrypt_flag_file_name='decrypt_data_done.flag'):
    """
    Decrypt data files under path of decrypt_dir_name
    Args:
        dew_config_file: the config file of DEW service
        decrypt_dir_name: directory name of encrypted data file
        user_dir_prefix: user directory prefix
        decrypt_flag_file_name: decrypt flag file name
    Returns:
        path of plain data file
    """
    decrypt_flag_file = os.path.join(user_dir_prefix, decrypt_flag_file_name)
    if int(CUR_RANK_ID) == 0 and not os.path.exists(decrypt_dir_name):
        raise DewError(
            f"[{CUR_RANK_ID}] data cipher dir : {decrypt_dir_name} not exists")
    if decrypt_dir_name[-1] == '/':
        plain_dir = os.path.join(os.path.dirname(decrypt_dir_name[:-1]),
                                 'data_plain')
    else:
        plain_dir = os.path.join(os.path.dirname(decrypt_dir_name),
                                 'data_plain')
    if int(CUR_RANK_ID) != 0:
        while not os.path.exists(decrypt_flag_file):
            time.sleep(5)
            print(f"[{CUR_RANK_ID}] waiting for rank id 0 to decrypt data...")
        return plain_dir

    dew_config = get_dew_config(dew_config_file)
    dew_inst = Dew(
        base_url=dew_config.get('base_url'),
        project_id=dew_config.get('project_id'),
        key_id=dew_config.get('key_id'),
        ak=dew_config.get('ak'),
        sk=dew_config.get('sk'),
        proxies=dew_config.get('proxies') if 'proxies' in dew_config else None,
        verify=dew_config.get('verify')
    )

    # 获取数据密钥的明文
    print(f"[{CUR_RANK_ID}] start to get data key from dew service")
    cipher_key = dew_config.get('cipher_key')
    plain_key = dew_inst.get_data_key(cipher_key)
    plain_key = hex_string_to_bytes(plain_key)
    print(f"[{CUR_RANK_ID}] get data key from dew service successfully")

    if not os.path.exists(plain_dir):
        os.makedirs(plain_dir)
    # 遍历数据集目录，解密数据集
    for path, dir_list, file_list in os.walk(decrypt_dir_name):
        for file_name in file_list:
            encrypted_data_file = os.path.join(path, file_name)
            dst_dir = path.replace(os.path.realpath(decrypt_dir_name),
                                   os.path.realpath(plain_dir))
            decrypt_file_stream(encrypted_data_file, dst_dir, plain_key)
    with open(decrypt_flag_file, 'w') as f:
        f.write('ok')
    print(f'[{CUR_RANK_ID}] decrypt_dir done')
    return plain_dir


def delete_plain_dir(plain_dir):
    if plain_dir and os.path.exists(plain_dir) and int(CUR_RANK_ID) == 0:
        shutil.rmtree(plain_dir)
        print(f'[{CUR_RANK_ID}] remove dir : {plain_dir} successfully')
    else:
        print(f'[{CUR_RANK_ID}] no need remove')


def save_cipher_key_file(save_file, cipher_key):
    """
        将数据密钥的密文保存至dew_config_file中
        Args:
            save_file: 保存文件路径
            cipher_key: 数据密钥密文
    """
    with open(save_file, 'wb') as fd:
        fd.write(cipher_key.encode())
        print(f'[{CUR_RANK_ID}] save cipher key successfully : {save_file}')


def get_cipher_and_signature_id(crypto_id_str, signature_id_str):
    crypto_algo = (
        CryptoAlgorithm.AES128_CBC,
        CryptoAlgorithm.AES256_CBC,
        CryptoAlgorithm.AES128_GCM,
        CryptoAlgorithm.AES256_GCM
    )
    sign_algo = (
        SignAlgorithm.HMAC_SHA256,
        SignAlgorithm.HMAC_SHA384,
        SignAlgorithm.HMAC_SHA512,
    )
    return crypto_algo[int(crypto_id_str) - 1], sign_algo[
        int(signature_id_str) - 1]


def encrypt_file_stream(plain_key, src_file, dst_file, crypto_alg_id,
                        hmac_alg_id, read_batch_size):
    """
    Encrypt src_file and save result to dst_dir
    Args:
        plain_key: the plain data key
        src_file: plain file to be decrypted
        dst_file: cipher file
        crypto_alg_id: crypto algorithm id
        hmac_alg_id: signature algorithm id
        read_batch_size: read size of each time
    Returns:
        None
    """
    enc_stream = symmetric.encrypt_stream(plain_key,
                                          src_file,
                                          read_batch_size,
                                          crypto_alg_id,
                                          hmac_alg_id)
    with open(dst_file, 'wb') as f:
        for enc_data in enc_stream:
            f.write(enc_data)


def is_need_encrypt_file(file_name):
    return file_name.endswith('.ckpt') or file_name.endswith(
        '.air') or file_name.endswith('-graph.meta')


def get_data_key_from_dew(cipher_key, dew, dst_dir):
    if cipher_key:
        # 配置文件已经配置数据密钥密文，则调用dew解密
        plain_key = dew.get_data_key(cipher_key)
        print(f'[{CUR_RANK_ID}] get data key from DEW service successfully')
        return plain_key
    else:
        # 需要创建数据密钥
        res = dew.create_data_key()
        plain_key = res.get('plain_text')
        cipher_key = res.get('cipher_text')
        save_cipher_key_file(
            os.path.join(os.path.dirname(dst_dir), 'cipher_data_key.key'),
            cipher_key)
        return plain_key


def encrypt_model_dir(dew_config_file, model_dir, remove_plain_model=True):
    """
    对模型目录下的ckpt模型文件加密
    Args:
        dew_config_file: DEW配置文件
        model_dir: 训练生成的模型文件所在目录
        remove_plain_model: 对明文的模型文件加密后，是否删除明文模型文件，
            默认True为删除，False为保留
    """
    dew_config = get_dew_config(dew_config_file)
    dew_inst = Dew(
        base_url=dew_config.get('base_url'),
        project_id=dew_config.get('project_id'),
        key_id=dew_config.get('key_id'),
        ak=dew_config.get('ak'),
        sk=dew_config.get('sk'),
        proxies=dew_config.get('proxies') if 'proxies' in dew_config else None,
        verify=dew_config.get('verify')
    )
    # 首先从配置文件获取，如果不存在，则从DEW服务创建一个数据密钥对模型加密，数据密钥保存到配置文件中
    cipher_key = dew_config.get('cipher_key')
    plain_key = get_data_key_from_dew(cipher_key,
                                      dew_inst,
                                      model_dir)
    if not plain_key:
        raise DewError(f'[{CUR_RANK_ID}] plain data key invalid')
    plain_key = hex_string_to_bytes(plain_key)
    cipher_model_dir = os.path.join(os.path.dirname(model_dir), 'model_cipher')
    print(f'[{CUR_RANK_ID}] cipher model dir : {cipher_model_dir}')
    if not os.path.exists(cipher_model_dir):
        os.makedirs(cipher_model_dir)
    else:
        shutil.rmtree(cipher_model_dir)

    crypto_alg_id, hmac_alg_id = get_cipher_and_signature_id(
        dew_config.get('crypto'),
        dew_config.get('signature'))
    # 遍历模型目录，加密模型文件
    for path, dir_list, file_list in os.walk(model_dir):
        for file_name in file_list:
            if not is_need_encrypt_file(file_name):
                continue
            plain_model_file = os.path.join(path, file_name)
            cipher_model_file = os.path.join(cipher_model_dir, file_name)
            print(f'[{CUR_RANK_ID}] plain_model_file: {plain_model_file}')
            print(f'[{CUR_RANK_ID}] cipher_model_file : {cipher_model_file}')
            encrypt_file_stream(plain_key, plain_model_file, cipher_model_file,
                                crypto_alg_id, hmac_alg_id,
                                symmetric.CHUNK_SIZE_1_MB)
            print(f'[{CUR_RANK_ID}] encrypt {plain_model_file} successfully')
            if remove_plain_model:
                os.remove(plain_model_file)
                print(f'[{CUR_RANK_ID}] remove {plain_model_file} successfully')
    print(f'[{CUR_RANK_ID}] encrypt model dir done')


def encrypt_model(dew_config_file, model, des_dir):
    """
    对模型文件加密
    Args:
        dew_config_file: DEW配置文件
        model: 内存中的模型文件
        des_file_name: 模型密文文件路径
    """
    dew_config = get_dew_config(dew_config_file)
    dew_inst = Dew(
        base_url=dew_config.get('base_url'),
        project_id=dew_config.get('project_id'),
        key_id=dew_config.get('key_id'),
        ak=dew_config.get('ak'),
        sk=dew_config.get('sk'),
        proxies=dew_config.get('proxies') if 'proxies' in dew_config else None,
        verify=dew_config.get('verify')
    )
    # 首先从配置文件获取，如果不存在，则从DEW服务创建一个数据密钥对模型加密，数据密钥保存到配置文件中
    cipher_key = dew_config.get('cipher_key')
    os.makedirs(des_dir, exist_ok=True)
    plain_key = get_data_key_from_dew(cipher_key, dew_inst, des_dir)
    if not plain_key:
        raise DewError(f'[{CUR_RANK_ID}] plain data key invalid')
    plain_key = hex_string_to_bytes(plain_key)
    # 加密保存模型结构
    f_obj = BytesIO()
    f_obj.write(model.to_json().encode('utf-8'))
    struct_cipher = os.path.join(des_dir, "model_struct_cipher.json")
    with open(struct_cipher, 'wb') as f:
        for i in symmetric.encrypt_stream(plain_key, f_obj):
            f.write(i)

    # 加密保存模型权重
    f_obj = BytesIO()
    weights = model.get_weights()
    f_obj.write(pickle.dumps(weights))
    weights_cipher = os.path.join(des_dir, "model_weights_cipher.pkl")
    with open(weights_cipher, 'wb') as f:
        for i in symmetric.encrypt_stream(plain_key, f_obj):
            f.write(i)
    print(f'[{CUR_RANK_ID}] encrypt model to {des_dir} successfully')
